"""
Hi there, this is Hu Jian from CASIA. I am excited to release this module under
open source license. It contains some basic but really useful functions for
developing your own agent.
The implementation of path finding in the module is inspired by this awesome
blog which I strongly recommend you to take a look at.
  https://www.redblobgames.com/pathfinding/a-star/introduction.html
Feel free to add or modify anything you want.
Have fun!

NOTE: After this release, the use of `core.utils` is discouraged and will be
deprecated. To migrate your code, just change every appearance of
`from core.utils.map import Map`
inside you agent module to
`from .map import Map`.

P.S. If you dislike these lengthy comments as much as I do, delete them.
That's fine. :stuck_out_tongue_winking_eye:
"""
import heapq
import json
import os
import pickle
import random

import numpy


class Map:
    def __init__(self, map_id, data_dir: str):
        """
        Load basic map data, move cost data and see data.

        You could do a lot more funny stuff using these three kind of data, e.g.
        get the whole see matrix of a given position via
        `self.see[mode][row, col]` or dynamically modify `self.cost` according
        to the observation of obstacles to generate customized move path.

        :param scenario: int
        """
        # get corresponding map id

        # with open(os.path.join(data_dir, "scenario.csv")) as f:
        #     for line in f:
        #         attrs = line.strip().split(",")
        #         if attrs[0].strip('"') == str(scenario):
        #             map_id = int(attrs[1].strip('"'))
        #             #map_folder = f"{data_dir}/map/map_{map_id}"
        #             map_folder = f"{data_dir}/map/map_94"
        #             break
        map_folder = f"{data_dir}/map/map_{map_id}"

        with open(os.path.join(map_folder, "basic.json")) as f:
            self.basic = json.load(f)["map_data"]

        self.max_row = len(self.basic)
        self.max_col = len(self.basic[0])

        with open(os.path.join(map_folder, "cost.pickle"), "rb") as f:
            self.cost = pickle.load(f)

        self.see = numpy.load(os.path.join(map_folder, "see.npz"))["data"]

    def is_valid(self, pos):
        """
        Check if `pos` is inside the map.

        :param pos: int
        :return: bool
        """
        row, col = divmod(pos, 100)
        return 0 <= row < self.max_row and 0 <= col < self.max_col

    def get_map_data(self):
        """Not very useful. Kept for backward compatibility."""
        return self.basic

    def get_neighbors(self, pos):
        """
        Get neighbors of `pos` in six directions.

        :param pos: int
        :return: List[int]
        """
        if not self.is_valid(pos):
            return []
        row, col = divmod(pos, 100)
        return self.basic[row][col]["neighbors"]

    def can_see(self, pos1, pos2, mode):
        """
        Check if `pos1` can see `pos2` with given `mode`.

        `self.see[mode]` is a `numpy.ndarray` that supports multidimensional
        indexing. So you could get the whole see matrix of a given position via
        `self.see[mode][row, col]`.

        By bit-masking see matrices of different positions, you could easily get
        all desired positions. For example,
        `np.argwhere((self.see[0][8, 24] & self.see[0][24, 8]) == True)` returns
        all positions that can be seen from both `0824` and `2408`. Tweak the
        condition and you could create a lot more interesting stuff.

        :param pos1: int
        :param pos2: int
        :param mode: int
        :return: bool
        """
        if (
                not self.is_valid(pos1)
                or not self.is_valid(pos2)
                or not 0 <= mode < len(self.see)
        ):
            return False
        row1, col1 = divmod(pos1, 100)
        row2, col2 = divmod(pos2, 100)
        return self.see[mode][row1, col1, row2, col2]

    def get_distance(self, pos1, pos2):
        """
        Get distance between `pos1` and `pos2`.

        :param pos1: int
        :param pos2: int
        :return: int
        """
        if not self.is_valid(pos1) or not self.is_valid(pos2):
            return -1
        # convert position to cube coordinate
        row1, col1 = divmod(pos1, 100)
        q1 = col1 - (row1 - (row1 & 1)) // 2
        r1 = row1
        s1 = -q1 - r1
        row2, col2 = divmod(pos2, 100)
        q2 = col2 - (row2 - (row2 & 1)) // 2
        r2 = row2
        s2 = -q2 - r2
        # calculate Manhattan distance
        return (abs(q1 - q2) + abs(r1 - r2) + abs(s1 - s2)) // 2

    def gen_move_route(self, begin, end, mode):
        """
        Generate one of the fastest move path from `begin` to `end` with given
        `mode` using A* algorithm.

        Here I use python standard `heapq` package just for convenience. The
        idea of A* search algorithm is quite simple. It combines BFS with
        priority queue and adds some heuristic to accelerate the search.

        When an obstacle appears, you could modify the corresponding
        `self.cost[mode][row][col]` to some big number. Then this function is
        still able to give you one of the fastest path available.

        As you might have thought, it's ok to inject even more insights to the
        so-called `cost`. For example, when feeling a "threat", you could
        increase cost or add heuristic at some positions to avoid passing
        through these "dangerous" regions.

        For those who want to find all positions that can be reached within a
        given amount of time, Dijkstra is all you need. It's only BFS with
        priority queue. To find all positions that can be reached within a given
        amount of time, just modify the exit condition inside `search()` and
        save all valid positions during search. It's left for you brilliant
        developers to implement.

        :param begin: int
        :param end: int
        :param mode: int
        :return: List[int]
        """
        if (
                not self.is_valid(begin)
                or not self.is_valid(end)
                or not 0 <= mode < len(self.cost)
                or begin == end
        ):
            return []
        frontier = [(0, random.random(), begin)]
        cost_so_far = {begin: 0}
        came_from = {begin: None}

        def a_star_search():
            while frontier:
                _, _, cur = heapq.heappop(frontier)
                if cur == end:
                    break
                row, col = divmod(cur, 100)
                for neigh, edge_cost in self.cost[mode][row][col].items():
                    neigh_cost = cost_so_far[cur] + edge_cost
                    if neigh not in cost_so_far or neigh_cost < cost_so_far[neigh]:
                        cost_so_far[neigh] = neigh_cost
                        came_from[neigh] = cur
                        heuristic = self.get_distance(neigh, end)
                        heapq.heappush(frontier, (neigh_cost + heuristic, random.random(), neigh))

        def reconstruct_path():
            path = []
            if end in came_from:
                cur = end
                while cur != begin:
                    path.append(cur)
                    cur = came_from[cur]
                path.reverse()
            return path

        a_star_search()
        return reconstruct_path()

    def gen_move_route_ops(self, begin, end, mode, enemy_pos):
        """
        Generate one of the fastest move path from `begin` to `end` with given
        `mode` using A* algorithm.

        Here I use python standard `heapq` package just for convenience. The
        idea of A* search algorithm is quite simple. It combines BFS with
        priority queue and adds some heuristic to accelerate the search.

        When an obstacle appears, you could modify the corresponding
        `self.cost[mode][row][col]` to some big number. Then this function is
        still able to give you one of the fastest path available.

        As you might have thought, it's ok to inject even more insights to the
        so-called `cost`. For example, when feeling a "threat", you could
        increase cost or add heuristic at some positions to avoid passing
        through these "dangerous" regions.

        For those who want to find all positions that can be reached within a
        given amount of time, Dijkstra is all you need. It's only BFS with
        priority queue. To find all positions that can be reached within a given
        amount of time, just modify the exit condition inside `search()` and
        save all valid positions during search. It's left for you brilliant
        developers to implement.

        :param begin: int
        :param end: int
        :param mode: int
        :return: List[int]
        """
        if (
                not self.is_valid(begin)
                or not self.is_valid(end)
                or not 0 <= mode < len(self.cost)
                or begin == end
        ):
            return []
        frontier = [(0, random.random(), begin)]
        cost_so_far = {begin: 0}
        came_from = {begin: None}

        attention_pos = set()

        def a_star_search():
            while frontier:
                _, _, cur = heapq.heappop(frontier)
                if cur == end:
                    break
                row, col = divmod(cur, 100)
                for neigh, edge_cost in self.cost[mode][row][col].items():
                    enemy_cost = 0
                    if enemy_pos:
                        for pos in enemy_pos:
                            if self.can_see(pos, cur):
                                enemy_cost += 1e5
                                attention_pos.add(pos)

                    neigh_cost = cost_so_far[cur] + edge_cost + enemy_cost
                    if neigh not in cost_so_far or neigh_cost < cost_so_far[neigh]:
                        cost_so_far[neigh] = neigh_cost
                        came_from[neigh] = cur
                        heuristic = self.get_distance(neigh, end)
                        heapq.heappush(frontier, (neigh_cost + heuristic, random.random(), neigh))

        def reconstruct_path():
            path = []
            if end in came_from:
                cur = end
                while cur != begin:
                    path.append(cur)
                    cur = came_from[cur]
                path.reverse()
            return path

        a_star_search()
        return reconstruct_path(), list(attention_pos)

    def get_grid_type(self, pos):
        '''
        获取地形
        :param pos:
        :return: grid_type: 0-开阔地 1-丛林地 2-居民地 3-松软地 4-河流
        '''
        if not self.is_valid(pos):
            return 0
        return self.basic[pos // 100][pos % 100]["cond"]

    def is_in_cover(self, pos):
        """
        是否遮蔽地形
        :param pos:
        :return:
        """
        return self.get_grid_type(pos) in [1, 2]

    def is_in_soft(self, coord):
        if self.get_grid_type(coord) == 3:
            return True
        else:
            return False

    def get_grid_distance(self, source_grid: int, distance_start: int, distance_end: int):
        """
        计算在地图上source_grid的距离为distant内的范围
        :param source_grid: 基准点
        :param distance_start, distant_end: 距离
        :return: set
        """
        try:
            gridset = set()
            if distance_end == 0:
                gridset.add(source_grid)
                return gridset

            row = source_grid // 100
            col = source_grid % 100
            # 确定正方形的四个顶点
            row1 = min(row + distance_end + 1, self.max_row)
            row0 = max(row - distance_end, 0)
            col1 = min(col + distance_end + 1, self.max_col)
            col0 = max(col - distance_end, 0)

            for row_x in range(row0, row1):
                for col_x in range(col0, col1):
                    pos = row_x * 100 + col_x
                    if self.is_valid(pos):
                        dis = self.get_distance(source_grid, pos)
                        if (dis >= distance_start) and (dis <= distance_end):
                            gridset.add(pos)
            return gridset

        except Exception as e:
            raise e

    def get_height(self, pos):
        """
        获取指定位置的高程
        :param pos:
        :return:
        """
        if not self.is_valid(pos):
            return 0
        return self.basic[pos // 100][pos % 100]["elev"]

    def get_height_change_level(self, pos1, pos2):
        """
        获取两个位置的高差变化等级 (从pos1进入pos2的高差)
        :param pos1:
        :param pos2:
        :return:
        """
        return (self.get_height(pos2) - self.get_height(pos1)) // 20

    def check_2bop_see(self, bop_a, bop_b, mod=0):
        """
        检查算子a能否观察到b
        :param bop_a : 主动观察者
        :param bop_b : 目标
        :mod: a, b : 位置互易
        :param t_terrain
        :returns : True 可观察 / False 不可观察
        :key : 首先基于地形进行HexIcuIsOK通视判断;
               然后考虑两个影响因素, 地形遮蔽和算子掩蔽状态, 两个因素对当前观察距离都有减半的影响
        """
        if mod:
            bop_a, bop_b = bop_b, bop_a

        try:
            if not bop_a or not bop_b:
                return False
            flag_iou = False

            pos_a = bop_a["cur_hex"]
            pos_b = bop_b["cur_hex"]
            hex_distance = self.get_distance(pos_a, pos_b)

            if isinstance(bop_a['observe_distance'], dict):
                max_distance = bop_a["observe_distance"][bop_b["type"]]
            elif isinstance(bop_a['observe_distance'], list):
                max_distance = bop_a["observe_distance"][bop_b["sub_type"]]
            else:
                return False

            # if max_distance != 0: # B1用于判断对方是否是直升机 24是直升机
            #     max_distance += bop_b["B1"]
            # 如果算子之间的距离已经超过了最大距离则直接判断为不可观察
            if hex_distance > max_distance or max_distance == 0:
                flag_iou = False
            else:
                fly_a = 1 if bop_a["type"] == 3 else 0
                fly_b = 1 if bop_b["type"] == 3 else 0
                mod = 1 if fly_a == 1 and fly_b == 1 else 2 if fly_a == 1 and fly_b == 0 else 3 if fly_a == 0 and fly_b == 1 else 0
                if self.can_see(pos_a, pos_b, mod):
                    # 地形通视判断通过没有遮挡,开始依据目标状态判断是否可观察
                    distance_threshold = max_distance
                    # 目标处于遮蔽地形,观察能力减半, 暂定空中单位不受地形影响
                    flag_in_hide = False
                    if bop_b["type"] in [1, 2]:
                        if bop_b["move_state"] == 4:
                            flag_in_hide = True
                        if self.is_in_cover(pos_b):
                            distance_threshold = distance_threshold // 2
                        # 观察者观察目标（车辆）, 高差大于一个单位高程/或观察者为空中单位, 掩蔽无效
                        if flag_in_hide and bop_b["type"] == 2:
                            if bop_a["type"] == 3:
                                flag_in_hide = False
                            else:
                                if self.get_height_change_level(pos_a, pos_b) <= -1:
                                    flag_in_hide = False
                        distance_threshold = distance_threshold // 2 if flag_in_hide else distance_threshold
                    # 比较计算完成后的观察距离和实际距离
                    flag_iou = (hex_distance <= distance_threshold)
            return flag_iou
        except Exception as e:
            raise e

    def get_grid_in_cover(self, gridset: set):
        '''
        计算source_grids中在遮蔽地形内的坐标
        :param gridset 坐标集合
        :return grids_in_cover 坐标集合中位于遮蔽地形的坐标的集合
        '''
        try:
            grids_in_cover = [pos for pos in gridset if self.is_valid(pos) and self.is_in_cover(pos)]
            return set(grids_in_cover)
        except Exception as e:
            raise e

    def get_PM_scope(self, grid: int):
        '''
        get see range of PM
        '''
        return self.get_grid_distance(grid, 0, 2) - set(self.get_grid_in_cover(self.get_grid_distance(grid, 2, 2)))
